#!/usr/bin/env python3
"""
Population PBPK Modeling Script

This script simulates PBPK modeling for 100 individuals from Chinese population,
incorporating physiological variability and generating population pharmacokinetic curves.

Based on individual_pbpk_simple.py and individual_pbpk_modeling.py

Author: BlackCat@Cppo
Version: 1.0.0
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os

# Add src path to system path
# sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../src'))
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
print(f"Project root path: {project_root}")
sys.path.insert(0, project_root)

from src.absorption.absorption_module import AbsorptionModel
from src.distribution.distribution_module import DistributionModel
from src.elimination.elimination_module import EliminationModel


class PopulationPBPKSimulator:
    """Population PBPK Simulator for Chinese Population"""

    def __init__(self, n_subjects=100):
        """Initialize population simulator"""
        self.n_subjects = n_subjects
        self.time_points = np.linspace(0, 24, 10000)  # 24 hours (1 day), 48 time points
        self.subjects_data = []
        self.simulation_results = {}
        self.population_results = {}

        # Chinese population physiological parameters (mean ± CV)
        self.pop_params = {
            # Body weight (kg) - Chinese adult population
            'body_weight': {
                'mean': 65.0,
                'cv': 0.15,
                'distribution': 'lognormal'
            },

            # Age (years) - Chinese adult population
            'age': {
                'mean': 40.0,
                'cv': 0.25,
                'distribution': 'normal',
                'min': 18,
                'max': 80
            },

            # Height (cm) - Chinese adult population
            'height': {
                'mean': 168.0,
                'cv': 0.08,
                'distribution': 'normal',
                'min': 140,
                'max': 190
            },

            # Physiological parameters with variability
            'Vp': {
                'mean': 3.0,
                'cv': 0.20,
                'distribution': 'lognormal'
            },  # Plasma volume (L)

            'Ve': {
                'mean': 12.0,
                'cv': 0.15,
                'distribution': 'lognormal'
            },  # Extracellular fluid volume (L)

            'Vt': {
                'mean': 12.0,
                'cv': 0.20,
                'distribution': 'lognormal'
            },  # Tissue volume (L)

            # Liver blood flow (L/h) - Chinese population adjusted
            'Qliver': {
                'mean': 90.0,
                'cv': 0.15,
                'distribution': 'lognormal'
            },

            # Renal blood flow (L/h)
            'Qrenal': {
                'mean': 65.0,
                'cv': 0.15,
                'distribution': 'lognormal'
            },

            # Enzyme activity variability (CYP3A4 - common in Chinese population)
            'CYP3A4_activity': {
                'mean': 1.0,
                'cv': 0.30,
                'distribution': 'lognormal'
            }
        }

    def generate_individual_parameters(self, subject_id):
        """Generate individual-specific parameters based on Chinese population data"""
        # Basic physiological parameters
        bw = self._generate_parameter('body_weight')
        age = self._generate_parameter('age')
        height = self._generate_parameter('height')

        # Calculate BMI
        bmi = bw / ((height / 100) ** 2)

        # Adjust volumes based on body weight
        Vp = self._generate_parameter('Vp') * (bw / 65.0)
        Ve = self._generate_parameter('Ve') * (bw / 65.0)
        Vt = self._generate_parameter('Vt') * (bw / 65.0)

        # Adjust clearance based on body weight and enzyme activity
        cyp_activity = self._generate_parameter('CYP3A4_activity')
        Qliver = self._generate_parameter('Qliver') * (bw / 65.0) ** 0.75
        Qrenal = self._generate_parameter('Qrenal') * (bw / 65.0) ** 0.75

        # Individual parameters
        individual_params = {
            'subject_id': subject_id,
            'body_weight': bw,
            'age': age,
            'height': height,
            'bmi': bmi,

            # Dose and absorption
            'dose': 100.0,  # mg
            'ka': np.random.lognormal(np.log(1.0), 0.3),  # Absorption rate
            'kd': np.random.lognormal(np.log(1.0), 0.3),  # Dissolution rate
            'fa': np.random.beta(9, 1),  # Absorption fraction
            'fg': np.random.beta(9.5, 0.5),  # Gut metabolism
            'fh': np.random.beta(9, 1),  # Hepatic metabolism
            'ts': np.random.lognormal(np.log(1.0), 0.2),  # Stomach emptying time
            'tlag': np.random.lognormal(np.log(0.1), 0.5),  # Lag time

            # Distribution volumes
            'Vp': Vp,
            'Ve': Ve,
            'Vt': Vt,
            'fu': np.random.beta(2, 18),  # Fraction unbound (mean 0.1)
            'BP': 1.0,  # Blood-to-plasma ratio

            # Tissue partition coefficients (with variability)
            'Kp_liver': np.random.lognormal(np.log(2.0), 0.2),
            'Kp_kidney': np.random.lognormal(np.log(2.5), 0.2),
            'Kp_brain': np.random.lognormal(np.log(0.5), 0.3),
            'Kp_lung': np.random.lognormal(np.log(1.5), 0.2),
            'Kp_heart': np.random.lognormal(np.log(1.8), 0.2),
            'Kp_muscle': np.random.lognormal(np.log(1.2), 0.15),
            'Kp_adipose': np.random.lognormal(np.log(0.8), 0.25),

            # Clearance parameters (consistent with individual model)
            # Use absolute clearance values (L/h) rather than scaled
            'CL_hepatic': 2.0 * cyp_activity * (Qliver / 90.0),

            # Renal clearance scaled by renal blood flow
            'CL_renal': 1.5 * (Qrenal / 65.0),

            # Biliary clearance scaled by liver blood flow
            'CL_biliary': 0.5 * (Qliver / 90.0),

            # Other clearance scaled by body weight
            'CL_other': 0.2 * (bw / 65.0) ** 0.75
        }

        return individual_params

    def _generate_parameter(self, param_name):
        """Generate parameter value based on population distribution"""
        param_info = self.pop_params[param_name]
        mean = param_info['mean']
        cv = param_info['cv']
        distribution = param_info['distribution']

        if distribution == 'lognormal':
            # Log-normal distribution
            sigma = np.sqrt(np.log(1 + cv ** 2))
            mu = np.log(mean) - 0.5 * sigma ** 2
            value = np.random.lognormal(mu, sigma)
        elif distribution == 'normal':
            # Normal distribution
            sigma = mean * cv
            value = np.random.normal(mean, sigma)
            # Apply bounds if specified
            if 'min' in param_info:
                value = max(param_info['min'], value)
            if 'max' in param_info:
                value = min(param_info['max'], value)
        elif distribution == 'beta':
            # Beta distribution (for parameters between 0 and 1)
            alpha = mean * (1 - mean) / cv ** 2 - mean
            beta = alpha * (1 - mean) / mean
            value = np.random.beta(alpha, beta)
        else:
            value = mean

        return value

    def run_subject_simulation(self, subject_id, absorption_model='first_order',
                               distribution_model='minimal_pbpk', params=None):
        """Run PBPK simulation for a single subject"""
        if params is None:
            params = self.generate_individual_parameters(subject_id)

        # Add additional absorption parameters based on model type
        if absorption_model == 'adam':
            params.update({
                'use_adam': True,
                'solubility': 1.0,
                'permeability': 1e-4,
                'particle_size': 50.0,
                'density': 1.2,
                'molecular_weight': 300.0,
                'pKa': 7.4,
                'logP': 2.0
            })

        # 1. Absorption phase
        abs_model = AbsorptionModel(params)
        absorption_result = abs_model.calculate_absorption(params, self.time_points)

        # 2. Distribution phase - use the specified distribution model
        dist_model = DistributionModel(params)
        absorption_input = {
            'dose': params['dose'],
            'cumulative': absorption_result['cumulative_absorption'],
            'plasma_concentration': absorption_result['plasma_concentration'],
            'bioavailability': absorption_result['bioavailability']
        }

        if distribution_model == 'minimal_pbpk':
            # Simplified PBPK distribution model (一室模型)
            distribution_result = dist_model.calculate_distribution(
                params, self.time_points, absorption_input
            )
        elif distribution_model == 'full_pbpk':
            # Full PBPK distribution model (多室模型)
            multi_comp_result = dist_model.calculate_multi_compartment_plasma_concentration(
                params, self.time_points, absorption_input
            )

            # Calculate distribution volume
            tissue_partition = dist_model.calculate_tissue_partition_coefficients(params)
            distribution_volume = dist_model.calculate_V_d(params, tissue_partition)

            # Calculate brain distribution
            brain_result = dist_model.calculate_brain_distribution(
                params, self.time_points, multi_comp_result['plasma_concentration']
            )

            # Merge tissue concentrations
            tissue_concentrations = multi_comp_result.get('tissue_concentrations', {})
            tissue_concentrations['brain'] = brain_result['concentration']

            distribution_result = {
                'plasma_concentration': multi_comp_result['plasma_concentration'],
                'distribution_volume': distribution_volume,
                'tissue_concentrations': tissue_concentrations,
                'brain_concentration': brain_result['concentration'],
                'brain_partition': brain_result['partition_coefficient']
            }
        else:
            raise ValueError(f"Unknown distribution model type: {distribution_model}")

        # 3. Elimination phase
        elim_model = EliminationModel(params)

        # 对于完整PBPK模型，分布结果已经包含消除过程
        # 对于最小模型，需要计算消除
        if distribution_model == 'full_pbpk':
            # 完整PBPK模型，直接使用分布结果
            plasma_concentration = distribution_result['plasma_concentration']
        else:
            # 最小模型，需要计算消除
            elimination_result = elim_model.calculate_elimination(
                params, self.time_points, distribution_result
            )
            plasma_concentration = elimination_result['plasma_concentration']

        # Calculate PK metrics
        pk_metrics = self.calculate_pk_metrics(
            self.time_points, plasma_concentration
        )

        return {
            'subject_id': subject_id,
            'params': params,
            'time_points': self.time_points,
            'plasma_concentration': plasma_concentration,
            'pk_metrics': pk_metrics,
            'absorption': absorption_result,
            'distribution': distribution_result,
            'elimination': {'plasma_concentration': plasma_concentration}
        }

    def run_population_simulation(self, absorption_model='first_order',
                                  distribution_model='minimal_pbpk'):
        """Run simulation for entire population"""
        print(f"\nRunning population simulation for {self.n_subjects} subjects...")
        print(f"Model: {absorption_model} absorption + {distribution_model} distribution")

        results = []

        for subject_id in range(1, self.n_subjects + 1):
            try:
                result = self.run_subject_simulation(
                    subject_id, absorption_model, distribution_model
                )
                results.append(result)

                if subject_id % 20 == 0:
                    print(f"  Completed {subject_id}/{self.n_subjects} subjects")

            except Exception as e:
                print(f"  Error in subject {subject_id}: {str(e)}")
                continue

        # Store results
        key = f"{absorption_model}_{distribution_model}"
        self.population_results[key] = results

        print(f"  Successfully simulated {len(results)} subjects")

        return results

    def calculate_pk_metrics(self, time_points, concentration):
        """Calculate pharmacokinetic parameters"""
        # C_max and T_max
        C_max = np.max(concentration)
        T_max = time_points[np.argmax(concentration)]

        # AUC (trapezoidal rule)
        AUC = np.trapezoid(concentration, time_points)

        # Terminal half-life
        if len(concentration) >= 10:
            terminal_points = concentration[-10:]
            if np.any(terminal_points > 0):
                log_concentration = np.log(terminal_points[terminal_points > 0])
                terminal_times = time_points[-len(terminal_points):][terminal_points > 0]
                if len(log_concentration) > 1:
                    slope, _ = np.polyfit(terminal_times, log_concentration, 1)
                    t_half = -0.693 / slope if slope < 0 else 999
                else:
                    t_half = 999
            else:
                t_half = 999
        else:
            t_half = 999

        return {
            'C_max': C_max,
            'T_max': T_max,
            'AUC': AUC,
            't_half': t_half
        }

    def plot_population_curves(self, model_key=None):
        """Plot population pharmacokinetic curves"""
        if model_key is None:
            # Plot all models
            fig, axes = plt.subplots(2, 2, figsize=(15, 12))
            axes = axes.flatten()

            model_combinations = [
                ('first_order', 'minimal_pbpk'),
                ('first_order', 'full_pbpk'),
                ('adam', 'minimal_pbpk'),
                ('adam', 'full_pbpk')
            ]

            for idx, (abs_model, dist_model) in enumerate(model_combinations):
                key = f"{abs_model}_{dist_model}"
                ax = axes[idx]

                if key in self.population_results:
                    results = self.population_results[key]
                    self._plot_single_population(ax, results, key)
                else:
                    ax.text(0.5, 0.5, 'No Data', ha='center', va='center', fontsize=12)

            plt.tight_layout()
            plt.savefig('../output/population_pbpk_all_models.png', dpi=300, bbox_inches='tight')
            plt.show()

        else:
            # Plot specific model
            if model_key in self.population_results:
                fig, ax = plt.subplots(figsize=(10, 8))
                results = self.population_results[model_key]
                self._plot_single_population(ax, results, model_key)
                plt.tight_layout()
                plt.savefig(f'../output/population_pbpk_{model_key}.png', dpi=300, bbox_inches='tight')
                plt.show()

    def _plot_single_population(self, ax, results, model_key):
        """Plot single population model results"""
        # Plot individual curves
        for result in results[::5]:  # Plot every 5th subject for clarity
            ax.plot(result['time_points'], result['plasma_concentration'],
                    alpha=0.3, linewidth=0.5, color='lightblue')

        # Calculate and plot median
        all_concentrations = np.array([r['plasma_concentration'] for r in results])
        median_conc = np.median(all_concentrations, axis=0)
        ax.plot(results[0]['time_points'], median_conc, 'b-', linewidth=3, label='Median')

        # Calculate and plot percentiles
        p5 = np.percentile(all_concentrations, 5, axis=0)
        p95 = np.percentile(all_concentrations, 95, axis=0)
        ax.fill_between(results[0]['time_points'], p5, p95, alpha=0.3, color='blue',
                        label='5th-95th percentile')

        ax.set_xlabel('Time (h)')
        ax.set_ylabel('Plasma Concentration (mg/L)')
        ax.set_title(f'Population PK Curves\n{model_key.replace("_", " + ").title()}')
        ax.grid(True, alpha=0.3)
        ax.legend()

        # Add population statistics
        all_C_max = [r['pk_metrics']['C_max'] for r in results]
        all_auc = [r['pk_metrics']['AUC'] for r in results]

        stats_text = f'N = {len(results)}\n'
        stats_text += f'Median C_max = {np.median(all_C_max):.3f} mg/L\n'
        stats_text += f'Median AUC = {np.median(all_auc):.1f} mg·h/L\n'
        stats_text += f'CV(C_max) = {np.std(all_C_max) / np.mean(all_C_max) * 100:.1f}%\n'
        stats_text += f'CV(AUC) = {np.std(all_auc) / np.mean(all_auc) * 100:.1f}%'

        ax.text(0.02, 0.98, stats_text, transform=ax.transAxes,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

    def generate_summary_statistics(self):
        """Generate population PK parameter summary"""
        print("\n" + "=" * 80)
        print("Population Pharmacokinetic Parameters Summary")
        print("=" * 80)

        for model_key, results in self.population_results.items():
            print(f"\nModel: {model_key.replace('_', ' + ').title()}")
            print("-" * 50)

            # Extract PK parameters
            C_max_values = [r['pk_metrics']['C_max'] for r in results]
            auc_values = [r['pk_metrics']['AUC'] for r in results]
            tmax_values = [r['pk_metrics']['T_max'] for r in results]
            t_half_values = [r['pk_metrics']['t_half'] for r in results if r['pk_metrics']['t_half'] < 100]

            # Calculate statistics
            for param_name, values in [('C_max (mg/L)', C_max_values),
                                       ('AUC (mg·h/L)', auc_values),
                                       ('T_max (h)', tmax_values)]:
                print(f"\n{param_name}:")
                print(f"  Mean ± SD: {np.mean(values):.3f} ± {np.std(values):.3f}")
                print(f"  Median: {np.median(values):.3f}")
                print(f"  Range: {np.min(values):.3f} - {np.max(values):.3f}")
                print(f"  CV (%): {np.std(values) / np.mean(values) * 100:.1f}")
                print(f"  5th-95th percentile: {np.percentile(values, 5):.3f} - {np.percentile(values, 95):.3f}")

            if t_half_values:
                print(f"\nHalf-life (h):")
                print(f"  Mean ± SD: {np.mean(t_half_values):.2f} ± {np.std(t_half_values):.2f}")
                print(f"  Median: {np.median(t_half_values):.2f}")
                print(f"  CV (%): {np.std(t_half_values) / np.mean(t_half_values) * 100:.1f}")

            # Body weight statistics
            bw_values = [r['params']['body_weight'] for r in results]
            print(f"\nBody Weight (kg):")
            print(f"  Mean ± SD: {np.mean(bw_values):.1f} ± {np.std(bw_values):.1f}")
            print(f"  Range: {np.min(bw_values):.1f} - {np.max(bw_values):.1f}")


def main():
    """Main function"""
    print("=" * 80)
    print("Population PBPK Modeling System - Chinese Population")
    print("=" * 80)

    # Create population simulator
    simulator = PopulationPBPKSimulator(n_subjects=100)

    # Define model combinations to run
    model_combinations = [
        ('first_order', 'minimal_pbpk'),
        ('first_order', 'full_pbpk'),
        ('adam', 'minimal_pbpk'),
        ('adam', 'full_pbpk')
    ]

    # Run simulations for all models
    for abs_model, dist_model in model_combinations:
        results = simulator.run_population_simulation(abs_model, dist_model)

    # Generate plots
    print("\nGenerating population plots...")
    simulator.plot_population_curves()

    # Generate summary statistics
    simulator.generate_summary_statistics()

    print("\n" + "=" * 80)
    print("Population simulation completed!")
    print("Plots saved in output/ directory:")
    print("  - population_pbpk_all_models.png")
    print("=" * 80)


if __name__ == "__main__":
    main()